#include "Huffman.h"
#include "Huffman_private.h"

#include <string.h>
#include <stdlib.h>

#ifdef WINDOWS_MAKE

void SetNodeCode(CHuffmanNode* pNode)
{
	CHuffmanNode* pParent = pNode->pParent;
	while(ISVALIDPTR(pParent) && pParent->nCodeLength)
	{
		pNode->dwCode <<= 1;
		pNode->dwCode |= pParent->dwCode;
		pNode->nCodeLength++;
		pParent = pParent->pParent;
	}
}

int ConstructHuffmanTree(CHuffmanNode nodes[], BOOL bSetCodes)
{
	HuffmanTree hufftree;
	CHuffmanNode* chuffnodes;
	CHuffmanNode* pNodes[511] = {0}, *pNode;
	// add used ascii to Huffman queue
	int nNodeCount = 0;
	int nCount, i;
	int nParentNode, nBackNode;

	for(nCount = 0; nCount < 256 && nodes[nCount].nFrequency; nCount++)
		pNodes[nNodeCount++] = &nodes[nCount];
	nParentNode = nNodeCount;
	nBackNode = nNodeCount-1;
	while(nBackNode > 0)
	{
		// parent node
		pNode = &nodes[nParentNode++];
		// pop first child
		pNode->pLeft = PopNode(pNodes, nBackNode--, FALSE);
		// pop second child
		pNode->pRight = PopNode(pNodes, nBackNode--, TRUE);
		// adjust parent of the two poped nodes
		pNode->pLeft->pParent = pNode->pRight->pParent = pNode;
		// adjust parent frequency
		pNode->nFrequency = pNode->pLeft->nFrequency + pNode->pRight->nFrequency;
		// insert parent node depending on its frequency
		for(i = nBackNode; i >= 0; i--)
		{
			if(pNodes[i]->nFrequency >= pNode->nFrequency)
				break;
		}
		memmove(pNodes+i+2, pNodes+i+1, (nParentNode-i-2)*sizeof(CHuffmanNode*));
		pNodes[i+1] = pNode;
		nBackNode++;
	}
	// Convert CHuffmanNodes tree to HuffmanTree
	qsort(pNodes, nParentNode, sizeof(CHuffmanNode*), asciiCompareForPointerArray);
	chuffnodes = (CHuffmanNode*)malloc(sizeof(CHuffmanNode) * nParentNode);
	for(i=0;i < nParentNode; i++)
	{
		chuffnodes[i].byAscii = pNodes[i]->byAscii;
		chuffnodes[i].dwCode = pNodes[i]->dwCode;
		chuffnodes[i].nCodeLength = pNodes[i]->nCodeLength;
		chuffnodes[i].nFrequency = pNodes[i]->nFrequency;
		chuffnodes[i].pLeft = pNodes[i]->pLeft;
		chuffnodes[i].pParent = pNodes[i]->pParent;
		chuffnodes[i].pRight = pNodes[i]->pRight;
	}
	hufftree = ConvertCHuffmanNodes2HuffmanTree(chuffnodes, nParentNode);
	// set tree leaves nodes code
	if(bSetCodes == TRUE)
		for(nCount = 0; nCount < nNodeCount; nCount++)
			SetNodeCode(&nodes[nCount]);

	return nNodeCount;
}

HuffmanTree ConvertCHuffmanNodes2HuffmanTree(CHuffmanNode* chuffnodes, int count)
{
	HuffmanTree tree;
	int idx, left_idx, right_idx, parent_idx;

	tree.count = count;
	tree.nodes = (HuffmanNode*)malloc((sizeof(HuffmanNode) - 2 * sizeof(char)) * tree.count);
	while(count > 0)
	{
		AccessHuffmanNode(tree.nodes, count)->byAscii = chuffnodes[count].byAscii;
		AccessHuffmanNode(tree.nodes, count)->nFrequency = chuffnodes[count].nFrequency;
		for(idx = 0;idx < tree.count; idx++)
		{
			if (chuffnodes[count].pLeft = &chuffnodes[idx]) left_idx = idx;
			else if (chuffnodes[count].pParent = &chuffnodes[idx]) parent_idx = idx;
			else if (chuffnodes[count].pRight = &chuffnodes[idx]) right_idx = idx;
		}
	}	
}

BOOL CompressHuffman(unsigned char *pSrc, int nSrcLen, unsigned char **pDes, int *nDesLen)
{ // Generate CHuffmanNode Tree
	CHuffmanNode nodes[511];
	int nCount;
	int nNodeCount, nNodeSize;
	int nDesIndex;
	UBYTE *pDesPtr;

	// initialize nodes ascii
	for(nCount = 0; nCount < 256; nCount++)
	{
		nodes[nCount].byAscii = nCount;		
	}
	for(nCount = 0; nCount < 511; nCount++)
	{
		nodes[nCount].nFrequency = 0;
		nodes[nCount].nCodeLength = 0;
		nodes[nCount].dwCode = 0;
		nodes[nCount].pLeft = 0;
		nodes[nCount].pRight = 0;
		nodes[nCount].pParent = 0;
	}

	// get ascii frequencies
	for(nCount = 0; nCount < nSrcLen; nCount++)
		nodes[pSrc[nCount]].nFrequency++;
	// sort ascii chars depending on frequency
	qsort(nodes, 256, sizeof(CHuffmanNode), frequencyCompare);
	// construct Huffman tree
	nNodeCount = ConstructHuffmanTree(nodes, TRUE);
	// construct compressed buffer
	nNodeSize = sizeof(int)+sizeof(UBYTE);
	*nDesLen = nSrcLen + nNodeCount * nNodeSize;
	if ((*pDes = (UBYTE*)malloc(*nDesLen)) == 0)
		return FALSE;
	pDesPtr = *pDes;
	memset(pDesPtr, 0, *nDesLen);
	// save source buffer length at the first DWORD
	*(DWORD*)pDesPtr = nSrcLen;
	pDesPtr += sizeof(DWORD);
	// save Huffman tree leaves count-1 (as it may be 256)
	*pDesPtr = nNodeCount-1;
	pDesPtr += sizeof(UBYTE);
	// save Huffman tree used leaves nodes
	for(nCount = 0; nCount < nNodeCount; nCount++)
	{	// the array sorted on frequency so used nodes come first
		memcpy(pDesPtr, &nodes[nCount], nNodeSize);
		pDesPtr += nNodeSize;
	}
	// sort nodes depending on ascii to can index nodes with its ascii value
	qsort(nodes, 256, sizeof(CHuffmanNode), asciiCompare);

	nDesIndex = 0;
	// loop to write codes
	for(nCount = 0; nCount < nSrcLen; nCount++)
	{
		*(DWORD*)(pDesPtr+(nDesIndex>>3)) |= nodes[pSrc[nCount]].dwCode << (nDesIndex&7);
		nDesIndex += nodes[pSrc[nCount]].nCodeLength;
	}
	// update destination length
	*nDesLen = (pDesPtr - (*pDes))+(nDesIndex+7)/8;
	pDesPtr = (UBYTE*)malloc(*nDesLen);
	memmove(pDesPtr, *pDes, *nDesLen);
	free (*pDes);
	*pDes = pDesPtr;
	//*pDes = (UBYTE*)realloc(pDes, *nDesLen); /* Curious Error related Heap corruption */

	return TRUE;
}

BOOL CompressDifferentialHuffman(BYTE *pSrc, int nSrcLen, BYTE **pDes, int *nDesLen)
{
	int i, tmp1, tmp2;
	tmp1 = pSrc[0];
	for(i=1;i<nSrcLen;i++)// Calculate Differential
	{
		tmp2 = pSrc[i] - tmp1;
		tmp1 = pSrc[i];
		pSrc[i] = tmp2;		
	}
	return CompressHuffman(pSrc, nSrcLen, pDes, nDesLen);
}

BOOL CompressHuffmanUsingTree(UBYTE *pSrc, int nSrcLen, UBYTE **pDes, int *nDesLen)
{ // using HuffmanTree
	int nDesIndex = 0, nCount;
	UBYTE *pDesPtr = *pDes;

	*nDesLen = nSrcLen;
	if ((*pDes = (UBYTE*)malloc(*nDesLen)) == 0) return FALSE;
	pDesPtr = *pDes;
	memset(pDesPtr, 0, *nDesLen);
	// loop to write codes
	for(nCount = 0; nCount < nSrcLen; nCount++)
	{
		*(DWORD*)(pDesPtr+(nDesIndex>>3)) |= pRoot[pSrc[nCount]].dwCode << (nDesIndex&7);
		nDesIndex += pRoot[pSrc[nCount]].nCodeLength;
	}
	// update destination length
	*nDesLen = (pDesPtr - (*pDes))+(nDesIndex+7)/8;
	*pDes = (UBYTE*)realloc(*pDes, *nDesLen);

	return TRUE;
}
#endif